# -*- coding: utf-8 -*-
# !/usr/bin/env python
"""
-------------------------------------------------
   File Name：     model
   Description :   
   Author :       lth
   date：          2022/12/13
-------------------------------------------------
   Change Activity:
                   2022/12/13 17:17: create this script
-------------------------------------------------
"""
__author__ = 'lth'

import torch.nn.functional as F
from torch import nn
from torchvision import models


class ResnetEncoderDecoder(nn.Module):
    def __init__(self, class_num=26):
        super(ResnetEncoderDecoder, self).__init__()
        self.bn = nn.BatchNorm2d(64)
        resnet = models.resnet18(pretrained=False)
        self.conv = nn.Conv2d(1, 64, kernel_size=(3, 3), padding=(1, 1), stride=(1, 1))
        self.cnn = nn.Sequential(*list(resnet.children())[4:-2])
        self.out = nn.Linear(512, class_num + 1)

    def forward(self, input):
        input = F.relu(self.bn(self.conv(input)), True)
        input = F.max_pool2d(input, kernel_size=(2, 2), stride=(2, 2))
        input = self.cnn(input)

        input = input.permute(0, 2, 3, 1)
        input = F.softmax(self.out(input), dim=-1)
        # input [n,7,7,27]  27 represents the 26 + 1 ("ABCEDEFGHIJKLMNOPQRSTUVWXYZ_)
        # labels = labels.cuda()

        return input
